﻿// Copyright (c) Microsoft Open Technologies, Inc. All rights reserved. 
// See License.txt in the project root for license information.
using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Diagnostics.Contracts;
using System.Linq;
using System.Text;
using System.Threading.Tasks;

namespace System.Threading.Actors.Utilities
{
    // Provides multi-buffered state manager.
    public class StateBufferingHelper : IDisposable
    {
        // Outer dictionary needs to be ConcurrentDictionary because it may be updated/read concurrently by different threads.
        // Inner dictionary needs to be ConcurrentDictionary because it might be read by an actor method at the same time
        // that it is being updated by replication logic.
        private readonly ConcurrentDictionary<string, ConcurrentDictionary<string, object>> dictionary = new ConcurrentDictionary<string, ConcurrentDictionary<string, object>>(); // generic symbol table, durably stored

        // Three layers of operation caching:
        //      1. currentOperations maps a partition name to a PendingOperationsHolder that holds the current operations 
        //          for that partition name.  It exists to allow us to roll back the current delegate's operations in the 
        //          event that the delegate cannot complete.
        //      2. pendingOperations holds all operations waiting to be replicated.  The current ops
        //          are merged into here after the Actor.Eval delegate completes successfully.
        //      3. replicatingOperations holds operations that are currently being replicated.  This allows
        //          us to lock out modifications to the operations that are currently being replicated.
        //
        private ConcurrentDictionary<string, PendingOperationsHolder> currentOperations = new ConcurrentDictionary<string, PendingOperationsHolder>();
        private AggregateOperationsHolder pendingOperations = new AggregateOperationsHolder();
        private AggregateOperationsHolder replicatingOperations = new AggregateOperationsHolder();

        // Some general thoughts on locking:
        // The following happen under global lock:
        //      -- Snapshots taken in Get()/TryGet()/GetAllKeysWithPrefix()
        //      -- ReplicateAsync replacing replicatingOperations with pendingOperations and re-constructing
        //          a new pendingOperations.
        //      -- During Flush(), when an actor's current state is coalesced into pendingOperations.
        // This means that you might run into the following benign race condition:
        //      -- The durable store ("dictionary") could potentially be updated during the replication
        //          process concurrently with being read by Get/TryGet/GetAllKeysWithPrefix.  This is
        //          OK as the durable store is implemented with ConcurrentDictionary, which allows
        //          concurrent mutation and enumeration.



        // Signaling this wakes up the replicator task
        private AsyncAutoResetEvent replicationNeededEvent = new AsyncAutoResetEvent();

        private IPersistenceHandler m_parentState;

        private volatile bool m_shuttingDown = false;

        // Cache a pre-completed task for "early exit" situations so that we don't
        // have to create a new Task each time.
        private static Task<bool> s_preCompletedTask = Task.FromResult(true);

        public StateBufferingHelper(IPersistenceHandler parentState)
        {
            m_parentState = parentState;
            // Launch our async replication task loop
            ReplicationTaskLoop(); // fire-and-forget
        }

        public void Dispose()
        {
            if (replicationNeededEvent != null && !m_shuttingDown)
            {
                // This should shut down the replication task loop if it is currently waiting on
                // replicationNeededEvent
                m_shuttingDown = true;
                replicationNeededEvent.Set(); // This should shut down the replication task loop
                replicationNeededEvent.Dispose();
                replicationNeededEvent = null;
            }
        }

        public void Set(string partitionName, string key, object value)
        {
            Contract.Requires(partitionName != null);

            if (m_shuttingDown)
                throw new ObjectDisposedException("StateBufferingHelper");

            // Should be safe, as actor method should have exclusive access to currentOps.
            PendingOperationsHolder currentOps = currentOperations.GetOrAdd(partitionName, x => new PendingOperationsHolder());
            var newOp = ChangeOperation.FromSet(key, value);
            currentOps.ActionStore[key] = newOp;
        }

        public void Remove(string partitionName, string key)
        {
            Contract.Requires(partitionName != null);

            if (m_shuttingDown)
                throw new ObjectDisposedException("StateBufferingHelper");

            // Should be safe, as actor method should have exclusive access to currentOps.
            PendingOperationsHolder currentOps = currentOperations.GetOrAdd(partitionName, x => new PendingOperationsHolder());
            var newOp = ChangeOperation.FromRemove(key);
            currentOps.ActionStore[key] = newOp;
        }

        // Assumed to be called under state lock
        public bool TryGet(string partitionName, string key, out object obj)
        {
            if (m_shuttingDown)
                throw new ObjectDisposedException("StateBufferingHelper");

            obj = null;

            // A number of buffers that we will search for "key"
            PendingOperationsHolder currentOps = null;
            PendingOperationsHolder pendingOps = null;
            PendingOperationsHolder replicatingOps = null;
            ConcurrentDictionary<string, object> partitionNameDict = null;

            // This can be safely done outside of lock.  "currentOperations" is
            // a ConcurrentDictionary, to it's OK if it is read and written (by another thread)
            // concurrently.  And no other thread will be accessing the value for
            // the IsolatedState associated with partitionKey, because only
            // one such thread can be running at once (due to lock).
            currentOperations.TryGetValue(partitionName, out currentOps); 

            // Take a quick snapshot of the pending and replicating ops
            // Use lock to ensure that our snapshot is consistent.
            lock (m_parentState)
            {
                pendingOperations.ActionStore.TryGetValue(partitionName, out pendingOps);
                replicatingOperations.ActionStore.TryGetValue(partitionName, out replicatingOps);
                dictionary.TryGetValue(partitionName, out partitionNameDict);
            }

            // Attempt to find the desired key in the pending ops stack (current/pending/replicating)
            // THE ORDER OF CHECKING HERE IS IMPORTANT.  You want the newest changes to have precedence.
            ChangeOperation keyOperation;
            if (((currentOps != null) && currentOps.ActionStore.TryGetValue(key, out keyOperation)) ||
                 ((pendingOps != null) && pendingOps.ActionStore.TryGetValue(key, out keyOperation)) ||
                 ((replicatingOps != null) && replicatingOps.ActionStore.TryGetValue(key, out keyOperation)) )
            {
                if (keyOperation.OperationType == ChangeOperationType.Set)
                {
                    obj = keyOperation.OperationValue;
                    return true;
                }
                else if (keyOperation.OperationType == ChangeOperationType.Remove)
                {
                    return false;
                }
                String message = "Unrecognized CachedOperation subtype found in actor pending operations!";
                Contract.Assume(false, message);
                throw new InvalidOperationException(message);                
            }

            // Could not find key in pending ops stack.  Check the durable store.
            if(partitionNameDict != null && partitionNameDict.TryGetValue(key, out obj))
            {
                return true;
            }

            return false;
        }

        // Assumed to be called under state lock
        public object Get(string partitionName, string key) // TODO: Do we really need this?  
        {
            if (m_shuttingDown)
                throw new ObjectDisposedException("StateBufferingHelper");

            object result;
            if (!TryGet(partitionName, key, out result))
            {
                throw new KeyNotFoundException();
            }

            return result;
        }

        public Task Flush(string partitionName)
        {
            if (m_shuttingDown)
                throw new ObjectDisposedException("StateBufferingHelper");

            // Exit early if there are no current changes.
            PendingOperationsHolder currentChanges;
            if ((!currentOperations.TryGetValue(partitionName, out currentChanges)) || (currentChanges.ActionStore.Count == 0))
            {
                return s_preCompletedTask;
            }

            var tcs = new TaskCompletionSource<bool>();
            lock (m_parentState)
            {
                pendingOperations.Coalesce(currentChanges, partitionName, tcs);
                // Side-effect: Also replaces currentChanges._actionStore with new dictionary
            }

            replicationNeededEvent.Set(underLock: false);
            return tcs.Task;
        }

        // This loop is run by the replicatorTask
        private async void ReplicationTaskLoop()
        {
            while (true)
            {
                try
                {
                    // The reason that we use an AutoResetEvent here: If we have to manually
                    // reset the event, it's possible that we could "lose" a signal in a 
                    // scenario like the following:
                    //
                    //      Producer:                   Consumer:
                    //      ---------                   ---------
                    //                                  mre.WaitOne();
                    //      mre.Set();
                    //                                  (wake up, handle event)
                    //      mre.Set();
                    //                                  mre.Reset();
                    //                                  mre.WaitOne();
                    //                                  (potentially waiting forever)
                    //
                    // Now, to be fair, that only happens if you choose to reset the event
                    // right before waiting on it.  If you use the loop below and reset
                    // immediately after WaitOne() succeeds, then an MRE would probably be
                    // OK.
                    while (pendingOperations.ActionStore.Count == 0 && !m_shuttingDown)
                        await replicationNeededEvent.WaitAsync().ConfigureAwait(continueOnCapturedContext: false);
                }
                catch (ObjectDisposedException)
                {
                    // replicationNeededEvent has been disposed, probably as part of a cleanup operation.
                    // Fault any waiters and exit the loop.
                    ClearOperationsAndFaultWaiters(new ObjectDisposedException("Shutting down replication task loop - got an ObjectDisposedException while blocking."));
                    return;
                }

                //requestQueue.Take();
                //this.Logger.Write(Guid.Empty, "ReplicationTaskLoop woke up");

                // If we are being shut down, then break out of the loop, fault any waiters, and end this task
                if (m_shuttingDown)
                {
                    ClearOperationsAndFaultWaiters(new ObjectDisposedException("Shutting down replication task loop - someone triggered a shutdown (load balancing?)"));
                    return;
                }

                // All that we need to do under lock is to transition the PENDING ops into the REPLICATING ops.
                lock (m_parentState)
                {
                    this.replicatingOperations = pendingOperations;
                    pendingOperations = new AggregateOperationsHolder();
                }
                //this.Logger.Write(Guid.Empty, "ReplicationTaskLoop past lock");

                // Need to catch any errors thrown from implementation-specific replication/persistence.
                try
                {
                    await m_parentState.OnPersistAsync(replicatingOperations.ExportAllChanges()).ConfigureAwait(continueOnCapturedContext: false);                    
                    replicatingOperations.ReleaseWaitingTasks(true);
                    
                    // Don't do this -- it monkeys with the holding buffer outside of the lock, and
                    // can therefore confuse GetAllKeysWithPrefix.
                    // And don't null it out either.  It will be cleared/nulled/unrooted the next time
                    // that we run through this loop.
                    // TODO: We could consider putting it under lock...
                    //replicatingOperations.Clear();
                }
                catch (Exception e)
                {
                    var baseException = (e is AggregateException) ? ((AggregateException)e).GetBaseException() : e;
                    replicatingOperations.FaultWaitingTasks(baseException);
                    pendingOperations.FaultWaitingTasks(baseException);
                    lock (m_parentState)
                    {
                        replicatingOperations.Clear(); // These operations need to be considered lost/faulted, so let's not keep them around.
                        pendingOperations.Clear();
                    }
                }
            }
        }

        // These operations will be lost, and we need to fault all waiters.
        private void ClearOperationsAndFaultWaiters(Exception exception)
        {
            lock (m_parentState)
            {
                pendingOperations.FaultWaitingTasks(exception);
                pendingOperations.Clear();
                replicatingOperations.FaultWaitingTasks(exception);
                replicatingOperations.Clear();
            }
        }

        // Allows us to get a keys of a certain prefix, stripped of the prefix
        // This comes in handy for set operations, trying to find all members of a set.
        public IEnumerable<string> GetAllKeysWithPrefix(string partitionName, string keyPrefix)
        {
            Contract.Requires(keyPrefix != null);

            if (m_shuttingDown)
                throw new ObjectDisposedException("StateBufferingHelper");

            var prefixLen = keyPrefix.Length;

            ConcurrentDictionary<string, object> partitionNameDict;
            PendingOperationsHolder currentOps, pendingOps, replicatingOps;
            currentOperations.TryGetValue(partitionName, out currentOps); // Can be safely done outside of lock
            // Take a snapshot of pending/replicating/durable store, using a lock to ensure consistency.
            lock (this.m_parentState)
            {
                // These should be done from inside the global lock.
                dictionary.TryGetValue(partitionName, out partitionNameDict);
                pendingOperations.ActionStore.TryGetValue(partitionName, out pendingOps);
                replicatingOperations.ActionStore.TryGetValue(partitionName, out replicatingOps);
            }

            HashSet<string> results;

            // If this partition has an entry in the "durable store", query it
            // for candidates.
            if (partitionNameDict != null)
            {
                results = new HashSet<string>(
                    partitionNameDict.Select(kvp => kvp.Key).
                        Where(s => s.StartsWith(keyPrefix)).
                        Select(s => s.Substring(prefixLen)));
            }
            else results = new HashSet<string>();


            // Need to take pending operations into account
            // Go in reverse order so that the most recent operations take precedence
            foreach (var holder in new PendingOperationsHolder[] {replicatingOps, pendingOps, currentOps})
            {
                if (holder == null) continue;
                foreach (var kvp in holder.ActionStore)
                {
                    if (!kvp.Key.StartsWith(keyPrefix))
                        continue;
                    String shortName = kvp.Key.Substring(prefixLen);
                    if (kvp.Value.OperationType == ChangeOperationType.Set) results.Add(shortName);
                    else if (kvp.Value.OperationType == ChangeOperationType.Remove) results.Remove(shortName);
                    else throw new InvalidOperationException("Unknown CachedOperation type");
                }
            }

            return results;
        }

        // Allow direct access to durable store.
        // Provided so that end-of-the-line logic (like Windows Fabric replication logic) can store
        // permanent values, rather than writing to a transaction buffer.
        public ConcurrentDictionary<string,object> GetPermaStore(string partitionName)
        {
            Contract.Requires(partitionName != null);

            ConcurrentDictionary<string, object> result;
            lock (dictionary)
            {
                result = dictionary.GetOrAdd(partitionName, (x => new ConcurrentDictionary<string, object>()));
            }
            return result;
        }

        public ConcurrentDictionary<string, ConcurrentDictionary<string, object>> PermaStore
        {
            get { return dictionary; }
        }

        //public Dictionary<string, object> PermaStore[string partitionName]
        //{
        //    get { return dictionary; }
        //}

        public bool ChangesPendingForCurrentOperation(string partitionName)
        {
            PendingOperationsHolder currentOps;
            if (currentOperations.TryGetValue(partitionName, out currentOps))
            {
                return currentOps.ActionStore.Count > 0;
            }
            else return false;
        }

        public void ClearChangesForCurrentOperation(string partitionName)
        {
            PendingOperationsHolder currentOps;
            if (currentOperations.TryGetValue(partitionName, out currentOps))
            {
                currentOps.Clear();
            }
        }
    }

    // This class provides a generic holding space for pending operations (i.e., Set/Remove operations).
    // TODO: This is really just a Dictionary<string,ChangeOperation> now.  Is it still valuable to
    // have a PendingOperationsHolder class?
    public class PendingOperationsHolder
    {
        // A store for pending operations
        private Dictionary<string, ChangeOperation> _actionStore = new Dictionary<string, ChangeOperation>();

        public PendingOperationsHolder()
        {
        }

        // A special copy constructor used during Coalescing
        internal PendingOperationsHolder(PendingOperationsHolder other)
            : this()
        {
            this._actionStore = other._actionStore;
            other._actionStore = new Dictionary<string, ChangeOperation>();
        }

        // Clears out all pending operations and waiters upon those operations
        public void Clear()
        {
            _actionStore.Clear();
        }

        // Allow access to the action store
        internal Dictionary<string, ChangeOperation> ActionStore { get { return _actionStore; } }
    }

    // This class provides a generic holding space for pending operations (i.e., Set/Remove operations),
    // potentially across multiple IsolatedActorStates.  This class is used to hold both the "pending"
    // operations and the "replicating" operations.
    internal class AggregateOperationsHolder
    {
        // A store for pending operations
        private Dictionary<string, PendingOperationsHolder> _actionStore = new Dictionary<string, PendingOperationsHolder>();
        // Keep track of who is waiting for the completion of pending operations replication
        private List<TaskCompletionSource<bool>> _waitingTasks = new List<TaskCompletionSource<bool>>();

        // Clears out all pending operations and waiters upon those operations
        public void Clear()
        {
            _actionStore.Clear();
            _waitingTasks.Clear();
        }

        // Merge newer operations into this holder
        // ASSUMED TO BE DONE UNDER GLOBAL LOCK.
        public void Coalesce(PendingOperationsHolder newerOperations, string partitionName, TaskCompletionSource<bool> waiter)
        {
            PendingOperationsHolder holder;
            if (_actionStore.TryGetValue(partitionName, out holder))
            {
                // We already have pending operations for this partition name.  Absorb the new
                // ones and clear "newerOperations".
                foreach (var kvp in newerOperations.ActionStore) holder.ActionStore[kvp.Key] = kvp.Value;
                newerOperations.Clear();
            }
            else
            {
                // Special copy-constructor that has the side-effect of clearing "newerOperations"
                holder = new PendingOperationsHolder(newerOperations);
                _actionStore[partitionName] = holder;
            }

            _waitingTasks.Add(waiter);
        }

        // Called upon successful replication of operations
        public void ReleaseWaitingTasks(bool disableInlineContinuations)
        {
            if (_waitingTasks.Count == 0)
                return;

            // TaskCompletionSource.TrySetResult and other helpers will run continuations.  If we are
            // not careful, we can run into deadlocks.  For example, if we pre-loaded methods here and
            // our PreloadMethods method was synchronous & blocked until replication was complete,
            // the replication task itself could call ReleaseWaitingTasksd, which runs continuations.  
            // One of those continuations could run PreloadMethods.  This resulted in a single thread 
            // deadlocking itself.  Fixed by changing PreloadMethods to be async, but we can still fix
            // our behavior here to avoid similar future problems.
            if (disableInlineContinuations)
            {
                // _waitingTasks could be modified after this method returns.  Clone it for correctness.
                var clonedList = new List<TaskCompletionSource<bool>>(_waitingTasks);
                Task.Run(() =>
                {
                    foreach (var waiter in clonedList) waiter.TrySetResult(true);
                });
            }
            else
                foreach (var waiter in _waitingTasks) waiter.TrySetResult(true);
        }

        // Called when the replication of operations faults
        public void FaultWaitingTasks(Exception exception)
        {
            foreach (var waiter in _waitingTasks) waiter.TrySetException(exception);
        }

        // Allow access to the action store
        internal Dictionary<string, PendingOperationsHolder> ActionStore { get { return _actionStore; } }

        internal List<ChangeOperation> ExportAllChanges()
        {
            List<ChangeOperation> result = new List<ChangeOperation>();
            foreach (var kvp in _actionStore)
            {
                var partitionName = kvp.Key;
                foreach (var kvpInner in kvp.Value.ActionStore)
                {
                    var changeOp = kvpInner.Value;
                    changeOp.PartitionName = partitionName;
                    result.Add(changeOp);
                }
            }

            return result;
        }
    }

    public enum ChangeOperationType
    {
        Set,
        Remove
    }

    public class ChangeOperation
    {
        public ChangeOperationType OperationType { get; private set; }
        public string OperationKey { get; private set; }
        public object OperationValue { get; private set; }
        public string PartitionName { get; set; }  // might be set after construction

        private ChangeOperation(ChangeOperationType operationType, string operationKey, object operationValue, string partitionName = null)
        {
            OperationType = operationType;
            OperationKey = operationKey;
            OperationValue = operationValue;
            PartitionName = partitionName;
        }

        public static ChangeOperation FromSet(string key, object value, string partitionName = null)
        {
            return new ChangeOperation(ChangeOperationType.Set, key, value, partitionName);
        }

        public static ChangeOperation FromRemove(string key, string partitionName = null)
        {
            return new ChangeOperation(ChangeOperationType.Remove, key, null, partitionName);
        }

        public override string ToString()
        {
            return String.Format("ChangeOperation {0}: Key=\"{1}\", Value=\"{2}\", PartitionName={3}",
                OperationType, OperationKey, OperationValue, PartitionName);
        }
    }

    // IActorState implementations provide this method (OnPersist) to handle implementation-specific
    // logic for persisting/replicating operations.
    public interface IPersistenceHandler
    {
        Task OnPersistAsync(IList<ChangeOperation> toPersist);
    }
}
